from __future__ import print_function, annotations

from llvmlite import ir
from llvmlite.ir import Constant

from common.utils import log_and_raise
from spel4llvm.err_defs import ParseError
from spel4llvm.grammar_value_types import (
    TP_INT,
    TP_FLOAT,
    TP_BOOL,
    TP_STR,
    GrammarValueType,
    infer_type_from_py_value,
)
from spel4llvm.llvm_types import (
    LLVM_CHAR_TYPE,
    LLVM_INT_TYPE,
    LLVM_FLOAT_TYPE,
    LLVM_BOOL_TYPE,
    LLVM_STR_TYPE,
)


def get_linkage(user_visible):
    return "internal" if user_visible else "private"


def str_value_to_llvm_storage(mod: ir.Module, str_value: str, global_var_name: str):
    str_sz = len(str_value) + 1
    char_arr_type = ir.ArrayType(LLVM_CHAR_TYPE, str_sz)
    str_var = ir.GlobalVariable(mod, char_arr_type, name=global_var_name)
    str_var.global_constant = True
    str_var.unnamed_addr = True  # addr is insignificant
    str_var.align = 1  # const str align=1
    str_var.linkage = "private"
    str_var.initializer = Constant(
        char_arr_type, bytearray(str_value.encode("utf-8") + b"\0")
    )
    return str_var


def str_value_to_llvm_globalvar(
    mod: ir.Module, str_value: str, global_var_name: str, user_visible=True
):
    storage_var = str_value_to_llvm_storage(
        mod, str_value, global_var_name + ".storage"
    )

    str_var = ir.GlobalVariable(mod, LLVM_STR_TYPE, name=global_var_name)
    str_var.global_constant = not user_visible
    str_var.unnamed_addr = True  # addr is insignificant
    str_var.align = 1  # const str align=1
    str_var.linkage = get_linkage(user_visible)
    str_var.initializer = storage_var.gep([LLVM_INT_TYPE(0), LLVM_INT_TYPE(0)])
    return str_var


def list_value_to_llvm_storage(
    module,
    elem_llvm_type,
    py_list: list,
    storage_name,
    builder: ir.IRBuilder,
    gen_ctx,
):
    arr_sz = len(py_list)
    arr_type = ir.ArrayType(elem_llvm_type, arr_sz)
    arr_var = ir.GlobalVariable(module, arr_type, name=storage_name)
    arr_var.global_constant = False
    arr_var.unnamed_addr = True
    arr_var.align = None
    arr_var.linkage = "private"

    for i, py_elem in enumerate(py_list):
        ptr = builder.gep(
            arr_var,
            [LLVM_INT_TYPE(0), LLVM_INT_TYPE(i)],
        )
        elem_globalvar = py_value_to_llvm_global(
            module,
            py_elem,
            infer_type_from_py_value(py_elem),
            "%s.%d" % (storage_name, i),
            builder,
            gen_ctx,
            user_visible=False,
        )
        gval = builder.load(elem_globalvar)
        builder.store(gval, ptr)
    return arr_var


def list_value_to_llvm_globalvar(
    module: ir.Module,
    grammar_value_typ: GrammarValueType,
    py_list: list,
    list_name: str,
    builder: ir.IRBuilder,
    gen_ctx,
    user_visible=True,
):
    elem_llvm_type = grammar_value_typ.elem_type.to_llvm_type()
    # storage part
    storage_var = list_value_to_llvm_storage(
        module,
        elem_llvm_type,
        py_list,
        list_name + ".storage",
        builder,
        gen_ctx,
    )

    # ref part
    list_type = ir.LiteralStructType([LLVM_INT_TYPE, elem_llvm_type.as_pointer()])
    list_var = ir.GlobalVariable(module, list_type, name=list_name)
    list_var.global_constant = True
    list_var.unnamed_addr = True
    list_var.align = None
    list_var.linkage = get_linkage(user_visible)
    list_var.initializer = ir.Constant(
        list_type,
        [
            LLVM_INT_TYPE(len(py_list)),
            storage_var.gep([LLVM_INT_TYPE(0), LLVM_INT_TYPE(0)]),
        ],
    )
    return list_var


def numeric_value_to_llvm_globalvar(
    mod: ir.Module,
    numeric_value,
    llvm_numeric_typ,
    global_var_name: str,
    user_visible=True,
):
    int_global_var = ir.GlobalVariable(mod, llvm_numeric_typ, name=global_var_name)
    int_global_var.global_constant = False
    int_global_var.unnamed_addr = True  # addr is insignificant
    int_global_var.align = 1  # const str align=1
    int_global_var.linkage = get_linkage(user_visible)
    int_global_var.initializer = Constant(llvm_numeric_typ, numeric_value)
    return int_global_var


def py_value_to_llvm_global(
    mod: ir.Module,
    py_value,
    grammar_value_typ,
    global_var_name: str,
    builder,
    gen_ctx,
    user_visible=True,
):
    # python has str instead of char, so here we dont process TP_CHAR
    if grammar_value_typ == TP_INT:
        gvar = numeric_value_to_llvm_globalvar(
            mod, py_value, LLVM_INT_TYPE, global_var_name, user_visible
        )
    elif grammar_value_typ == TP_FLOAT:
        gvar = numeric_value_to_llvm_globalvar(
            mod, py_value, LLVM_FLOAT_TYPE, global_var_name, user_visible
        )
    elif grammar_value_typ == TP_BOOL:
        gvar = numeric_value_to_llvm_globalvar(
            mod, py_value, LLVM_BOOL_TYPE, global_var_name, user_visible
        )
    elif grammar_value_typ == TP_STR:
        gvar = str_value_to_llvm_globalvar(mod, py_value, global_var_name, user_visible)
    elif grammar_value_typ.is_list_type():
        gvar = list_value_to_llvm_globalvar(
            mod,
            grammar_value_typ,
            py_value,
            global_var_name,
            builder,
            gen_ctx,
            user_visible,
        )
    else:
        log_and_raise(
            ParseError,
            "unsupported sym type:%s of sym:%s" % (grammar_value_typ, global_var_name),
        )
    return gvar
